import cv2 as cv
import numpy as np
import time
import os
import imutils


def image_show(name, image, time, show_state=False):
    if False:
        rate = 6
        width, height = int(image.shape[1] / rate), int(image.shape[0] / rate)
        showImage = cv.resize(image, (width, height), interpolation=cv.INTER_AREA)
        cv.namedWindow(name, cv.WINDOW_AUTOSIZE)
        cv.imshow(name, showImage)
        cv.waitKey(0)


def find_all_board(image, width, height, filterThres, areaRate):
    T, board = cv.threshold(image, filterThres + 15, 255, cv.THRESH_BINARY)
    print("TotalBoard 阈值 = {}".format(T))
    image_show("Board", board, 1, True)
    cnts = cv.findContours(board.copy(), cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)
    largeContours = [cnt for cnt in cnts if cv.contourArea(cnt) > areaRate * width * height]
    centerPoints = [(int(cv.minAreaRect(contour)[0][0]), int(cv.minAreaRect(contour)[0][1])) for contour in largeContours]
    for index, center in enumerate(centerPoints):
        anotherCenter = centerPoints[(index + 1) % len(centerPoints)]
        cv.line(board, center, anotherCenter, 255, 20)
    cnts = cv.findContours(board.copy(), cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)
    cntFinalFinded = sorted(cnts, key=cv.contourArea, reverse=True)[0]
    box = cv.minAreaRect(cntFinalFinded)
    (cx, cy), (w, h), angle = box
    if angle > 70:
        angle = 90 - angle
        box = (cx, cy), (h, w), angle
    elif angle < -70:
        angle = 90 + angle
        box = (cx, cy), (h, w), angle

    rectPoints = np.int0(cv.boxPoints(box))
    boxMask = np.zeros(image.shape, dtype=np.uint8)
    cv.drawContours(boxMask, [rectPoints], -1, 255, -1)
    imageBlackAround = cv.bitwise_and(image, boxMask)
    return imageBlackAround, box


def get_rect_points_by_clockwise(rectPoints):
    """将矩形的四个顶点按照顺时针方向排列"""
    # 先按照坐标和进行排列,坐标和最小的肯定是左上角,最大的肯定是右下角.
    # 另外两个点,按x坐标大小进行排列,x大的是右上角,x小的是左下角.
    rectPoints = sorted(rectPoints, key=np.sum)
    leftTop = rectPoints[0]
    rightBottom = rectPoints[3]

    if rectPoints[1][0] > rectPoints[2][0]:
        rightTop = rectPoints[1]
        leftBottom = rectPoints[2]
    else:
        rightTop = rectPoints[2]
        leftBottom = rectPoints[1]

    return leftTop, rightTop, rightBottom, leftBottom


def   distance(p1, p2):
    """计算平面坐标系中两个点之间的距离"""
    assert len(p1) == 2
    assert len(p2) == 2
    dis = np.sqrt(np.square(p2[0] - p1[0]) + np.square(p2[1] - p1[1]))
    return round(dis, 2)


def check_box(cnt, width, height, boardColor, calcThres):
    assert len(cnt) != 1
    res = {"result": "fail", "upBox": None, "downBox": None, "info": "", "calcSize": None, "calcErr": None}

    try:
        widthMax = width + calcThres * MM2PIX
        widthMin = width - calcThres*MM2PIX*0.3
        heightMax = height + calcThres * MM2PIX
        heightMin = height - calcThres*MM2PIX*0.3

        box = cv.minAreaRect(cnt)
        rectPoints = np.int0(cv.boxPoints(box))

        leftTop, rightTop, rightBottom, leftBottom = get_rect_points_by_clockwise(rectPoints)
        widthFinded = distance(leftTop, rightTop)
        heightFinded = distance(leftTop, leftBottom)

        res["width"] = int(widthFinded / MM2PIX)
        res["height"] = int(heightFinded / MM2PIX)
        if widthFinded < widthMin:
            res["result"] = "small"
            res["info"] = "识别错误,宽度识别值过小,可能尺寸填写过大造成.识别结果: ({:.2f},{:.2f}),工单尺寸: ({:.2f},{:.2f}),阈值: ({:.2f},{:.2f})".format(
                widthFinded, heightFinded, width, height, widthMin, widthMax)
            return res
        elif widthFinded > widthMax:
            res["result"] = "big"
            res["info"] = "识别错误,宽度识别值过大,可能尺寸填写过小造成.识别结果: ({:.2f},{:.2f}),工单尺寸: ({:.2f},{:.2f}),阈值: ({:.2f},{:.2f})".format(
                widthFinded, heightFinded, width, height, widthMin, widthMax)
            return res
        elif heightFinded < heightMin:
            res["result"] = "small"
            res["info"] = "识别错误,高度识别值过小,可能尺寸填写过大造成.识别结果: ({:.2f},{:.2f}),工单尺寸: ({:.2f},{:.2f}),阈值: ({:.2f},{:.2f})".format(
                widthFinded, heightFinded, width, height, widthMin, widthMax)
            return res
        elif heightFinded > heightMax:
            res["result"] = "big"
            res["info"] = "识别错误,高度识别值过大,可能尺寸填写过小造成.识别结果: ({:.2f},{:.2f}),工单尺寸: ({:.2f},{:.2f}),阈值: ({:.2f},{:.2f})".format(
                widthFinded, heightFinded, width, height, widthMin, widthMax)
            return res
        else:
            res["result"] = "ok"
            res["box"] = [leftTop, rightTop, rightBottom, leftBottom]
            res["info"] = ""
            res["calcSize"] = (res["width"], res["height"])
            res["calcErr"] = res["width"] - width // MM2PIX + res["height"] - height // MM2PIX
            return res
    except Exception as e:
        res["result"] = "error"
        res["info"] = str(e)
        return res


def check_inner_board(innerBoard, width, height, calcThres, boardColor):
    """检查内板是否符合要求"""
    useImg = boardColor.copy()
    boardArea = width * height
    res = {"result": "fail", "upBox": None, "downBox": None, "info": "", "calcSize": None, "calcErr": None}
    cnts = cv.findContours(innerBoard.copy(), cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)
    cntsSorted = sorted(cnts, key=cv.contourArea, reverse=True)
    s1 = 0
    try:
        _, (w1, h1), _ = cv.minAreaRect(cntsSorted[0])
        s = w1 * h1
    except Exception as e:
        pass

    if len(cnts) == 0:
        res["result"] = "fail"
        res["info"] = "找内板的时候没有找到任何轮廓."
    elif 1.2 * boardArea > s > 0.8 * boardArea:
        rectBox = cv.minAreaRect(cntsSorted[0])
        rectBoxPoints = np.int0(cv.boxPoints(rectBox))
        cv.drawContours(useImg, [rectBoxPoints], -1, (0, 255, 0), 5)
        image_show("FindInnerBox", useImg, 0, True)
        res = check_box(cntsSorted[0], width, height, boardColor, calcThres)
        return res
    elif boardArea >= 1.2 * boardArea:
        # 找到的板过大,可能是找不到中缝引起的.这个时候就要画线,如果画线找到的板还过大.
        pass


def get_innder_board(board, jointLine, width, height, calcThres, boardPos):
    rectWidth = max(int(width / 3), int(150 * MM2PIX))
    rectHeight = int(200 * MM2PIX)
    boardColor = cv.cvtColor(board, cv.COLOR_GRAY2BGR)
    image_show("OneBoardOriginal", board, 0, True)
    blurred = cv.GaussianBlur(board, (0, 0), 5)
    scharrX = cv.Scharr(blurred, cv.CV_64F, 1, 0)
    scharrX = cv.convertScaleAbs(scharrX)
    scharrY = cv.Scharr(blurred, cv.CV_64F, 0, 1)
    scharrY = cv.convertScaleAbs(scharrY)
    scharrXY = cv.addWeighted(scharrX, 0.5, scharrY, 0.5, 0)
    image_show("ScharrXY", scharrXY, 0, True)

    img32 = np.float32(scharrXY)
    square1 = cv.blur(np.multiply(img32, img32), (17, 17))
    square2 = np.multiply(cv.blur(img32, (17, 17)), cv.blur(img32, (17, 17)))
    squareDif32 = np.subtract(square1, square2)
    squareDif = squareDif32.copy()
    np.clip(squareDif32, 0, 255, out=squareDif)
    squareDif = squareDif.astype("uint8")
    image_show("SquareDif", squareDif, 1, True)

    thresInit = 3
    offset = 0
    resList = []
    for calcTime in range(1, 20):
        contourImg = squareDif.copy()
        k1 = cv.getStructuringElement(cv.MORPH_RECT, (17, 17))
        for i in range(10+calcTime):
            contourImg = cv.dilate(contourImg, k1)
        for i in range(10+calcTime):
            contourImg = cv.erode(contourImg, k1)
        image_show("ScharrXYDilate", contourImg, 1, True)
        thresUsed = contourImg[rectHeight:int(contourImg.shape[0] - rectHeight), rectWidth:contourImg.shape[1] - rectWidth]
        image_show("ThresUsed", thresUsed, 0, True)
        print("FindInnerBoardThres T = {},thresOffset = {}".format(np.mean(thresUsed), calcTime))
        _, contourImgNew = cv.threshold(contourImg, np.mean(thresUsed) + thresInit + calcTime, 255, cv.THRESH_BINARY)
        contourImgNew = cv.subtract(contourImgNew, jointLine)
        image_show("InnerBoardThresDealed", contourImgNew, 1, True)

        # 5. 进行泛红填充,将所有的内板区域全部填充为150,内板全部处理成白色,内板之外全部处理为黑色
        mask = np.zeros((contourImgNew.shape[0] + 2, contourImgNew.shape[1] + 2), dtype=np.uint8)
        cv.floodFill(contourImgNew, mask, (contourImgNew.shape[1] // 2, contourImgNew.shape[0] // 4), 150)
        cv.floodFill(contourImgNew, mask, (contourImgNew.shape[1] // 2, contourImgNew.shape[0] * 3 // 4), 150)
        floodFill = np.zeros(contourImgNew.shape, dtype=np.uint8)
        floodFill[contourImgNew == 150] = 255
        image_show("FloodFill", floodFill, 1, True)

        # 6.进行去角以前的最后的腐蚀膨胀
        k2 = cv.getStructuringElement(cv.MORPH_RECT, (17, 17))
        for i in range(10+calcTime):
            floodFill = cv.erode(floodFill, k2)
        for i in range(10+calcTime):
            floodFill = cv.dilate(floodFill, k2)
        image_show("FloodFillErode", floodFill, 1, True)

        # 7.去掉四个角
        if boardPos == "up":
            cv.rectangle(floodFill, (0, 0), (rectWidth, rectHeight), 0, -1)
            cv.rectangle(floodFill, (floodFill.shape[1], 0), (floodFill.shape[1] - rectWidth, rectHeight), 0, -1)
        elif boardPos == "down":
            cv.rectangle(floodFill, (0, floodFill.shape[0]), (rectWidth, floodFill.shape[0] - rectHeight), 0, -1)
            cv.rectangle(floodFill, (floodFill.shape[1], floodFill.shape[0]), (floodFill.shape[1] - rectWidth, floodFill.shape[0] - rectHeight), 0,
                         -1)
        image_show("MoveAngle", floodFill, 0, True)

        # 8.去验证最后寻找到的内板轮廓
        res = check_inner_board(floodFill, width, height, calcThres, boardColor)
        print("第{}次的识别结果,识别阈值:{},res:{}".format(calcTime, thresInit + calcTime, res))
        if res["result"] == "ok":
            resList.append([res["result"],res["info"],res["box"],res["calcSize"],res["calcErr"]])
        elif res["result"] == "big":
            break
    if len(resList) >= 1:
        resSorted = sorted(resList,key=lambda x:x[3])
        print("resSorted:{}".format(resSorted))
        return resSorted[0]
    return None, None, None, None,None


def get_all_boxes(image, width, height, calcThres, jointNumber):
    if image is None:
        return False, "图像异常,请重拍", None, None, None, None, None, None
    if len(image.shape) == 3:
        image = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
    image_show("Before", image, 1, True)
    imageGray, boxMaster = find_all_board(image, width, height, np.mean(image[:150, :1100]), 0.3)
    (cx, cy), (w, h), angle = boxMaster
    widthHalf = int(w / 2)
    heightHalf = int(h / 2)
    yStart = max(int(cy - heightHalf), 0)
    yEnd = min(int(cy + heightHalf), image.shape[0])
    xStart = max(int(cx - widthHalf), 0)
    xEnd = min(int(cx + widthHalf), image.shape[1])
    M = cv.getRotationMatrix2D((cx, cy), angle, 1)
    totalBoard = cv.warpAffine(imageGray, M, (image.shape[1], image.shape[0]))[yStart:yEnd, xStart:xEnd]
    jointLineMask = np.zeros(image.shape, dtype=np.uint8)
    jointLineMask[:, 1106 - 10:1106 + 10] = 255
    jointLineMaskRoated = cv.warpAffine(jointLineMask, M, (image.shape[1], image.shape[0]))[yStart:yEnd, xStart:xEnd]
    image_show("TotalBoard", totalBoard, 0, True)
    totalBoardColor = cv.cvtColor(totalBoard, cv.COLOR_GRAY2BGR)
    if jointNumber == "2":
        upDownSlice = int(20 * MM2PIX)
        upBoard = totalBoard.copy()[:int(heightHalf + upDownSlice), :]
        downBoard = totalBoard.copy()[int(heightHalf - upDownSlice):, :]
        upJointLine = jointLineMaskRoated.copy()[:int(heightHalf + upDownSlice), :]
        downJointLine = jointLineMaskRoated.copy()[int(heightHalf - upDownSlice):, :]
        upRes, upInfo, upBox,upCalcSize, upErr = get_innder_board(upBoard, upJointLine, width, height, calcThres, "up")
        downRes, downInfo, downBox,upCalcSize, downErr = get_innder_board(downBoard, downJointLine, width, height, calcThres, "down")
        if upRes == "ok" and downRes == "ok":
            downBoxPoint = [(p[0], p[1] + heightHalf - upDownSlice) for p in downBox]
            cv.drawContours(totalBoardColor, [np.array(upBox), np.array(downBoxPoint)], -1, (0, 255, 0), 2)
            image_show("TwoBoxes", totalBoardColor, 0, True)
            return "ok",upInfo+downInfo,upBox,downBox,upCalcSize,upErr+downErr,totalBoard

def get_width_and_height(fileName):
    """根据文件名或者宽度和高度"""
    if fileName[-4:] != ".bmp":
        return None, None
    if "w_" not in fileName or "h_" not in fileName:
        return None, None

    widthIndex = fileName.index("w_")
    heightIndex = fileName.index("h_")
    width = int(fileName[widthIndex + 2:widthIndex + 6])
    height = int(fileName[heightIndex + 2:heightIndex + 6])
    return width, height


if __name__ == '__main__':
    tryBegin = 2
    tryEnd = 3
    filePath = r"D:\chongda_raw\can't_find_middle\188--2020-09-21_16-45-46--w_1989--h_1745.bmp"
    MM2PIX = 3.2038
    mtx = [[0, 1, 0], [1, 0, 0]]
    nowTime = time.strftime("%Y-%m-%d_%H-%M-%S", time.localtime())
    keepDir = filePath
    if os.path.isfile(filePath):
        keepDir, _ = os.path.split(filePath)

    # 保存成功结果的实际路径
    # 保存失败图片的实际路径
    keepResPath = os.path.join(keepDir, nowTime.split("_")[0] + "_res")
    keepFailPath = os.path.join(keepDir, nowTime.split("_")[0] + "_fail")

    # 如果两个目录都不存在,这创建
    if not os.path.exists(keepResPath):
        os.mkdir(keepResPath)
    if not os.path.exists(keepFailPath):
        os.mkdir(keepFailPath)

    calcThres = 12
    fileNames = []
    remainEdge = [1, 1, 1, 1]
    jointNumber = "2"
    noCut = False if any(remainEdge) else True
    topPix, leftPix, bottomPix, rightPix = [int(x * MM2PIX) for x in remainEdge]
    if os.path.isfile(filePath):
        filePath, fileName = os.path.split(filePath)
        fileNames.append(fileName)
    else:
        for root, dirs, files in os.walk(filePath):
            if root != filePath:  # 过滤掉不是要处理的图像的目录
                continue
            for file in files:
                fileNames.append(file)

    for index, fileName in enumerate(fileNames):
        filePathReal = os.path.join(filePath, fileName)
        width, height = get_width_and_height(fileName)
        if width is None or height is None:
            print("*" * 10 + " 第{} 张图 ({}) 图像名称错误,已经跳过".format(index + 1, fileName) + "*" * 10)
            continue
        print("*" * 10 + " 第{} 张图 ({})".format(index + 1, fileName) + "*" * 10)
        img = cv.imread(filePathReal, cv.IMREAD_GRAYSCALE)
        resOkList = []  # 识别成功的列表
        resSizeDifList = []  # 识别尺寸出问题的列表
        totalBoard = None
        for tryTimes in range(tryBegin, tryEnd):
            calcStart = time.time()
            res, info, upBox, downBox, boardSize, boardError, totalBoard = get_all_boxes(img, width, height, calcThres, jointNumber)
            calcTime = time.time() - calcStart
            print("算法计算耗时:{}秒".format(calcTime))
            if res == "ok":
                resOkList.append((boardError, upBox, downBox, info, boardSize))
                break
            elif res == "sizeDif":
                resSizeDifList.append((boardError, upBox, downBox, info, boardSize))

        print("识别成功的结果 = {},个数 = {}".format(resOkList, len(resOkList)))
        print("识别尺寸出问题的结果 = {},个数 = {}".format(resSizeDifList, len(resSizeDifList)))
        if len(resOkList) == 0:
            # 识别失败,将图片保存到失败的目录下面
            cv.imwrite(os.path.join(keepFailPath, fileName), img)
            continue
        resSorted = sorted(resOkList, key=lambda x: x[0])
        upBox, downBox, info, boardSize = resSorted[0][1:]

        assert upBox is not None
        assert downBox is not None
        resPath = os.path.join(keepResPath, fileName)
        cutLines, imgRes, cutDis, boardLength, angleList = get_lines_with_boxes(img, upBox, downBox, mtx, topPix, leftPix, bottomPix, rightPix, width,
                                                                                height, totalBoard)
        cv.imwrite(resPath, imgRes)
